#pragma once

#include <algorithm>
#include <limits>
#include <vector>

#include "../utils/rand.h"

namespace mdtx::online::algo
{
    /**
     * @brief KLL sketch updater
     *
     * @cite Karnin, Z., Lang, K., & Liberty, E. (2016, October). Optimal quantile approximation in streams. In 2016 ieee 57th annual symposium on foundations of computer science (focs) (pp. 71-78). IEEE.
     *
     * @tparam T base type
     * @tparam P prob type
     */
    template <class T, class P = double>
    class kll_sketch
    {
    private:
        struct tpair
        {
            T v;
            size_t w;
            friend bool operator<(const tpair &lhs, const tpair &rhs)
            {
                return lhs.v < rhs.v;
            }
        };

        using tvect = std::vector<tpair>;
        using tcomp = std::vector<tvect>;

        size_t k, size{}, max_size{};
        bool flip{};
        T maxval{std::numeric_limits<T>::lowest()};

        xorshift128_bool rbool{};

        tcomp compact{};

        mutable tvect work1, work2, work3, work4;
        mutable std::vector<T> workv{};
        mutable std::vector<P> workp{};

        static constexpr size_t pow3[33] = {
            1, 3, 9, 27, 81, 243, 729, 2187, 6561, 19683, 59049, 177147, 531441,
            1594323, 4782969, 14348907, 43046721, 129140163, 387420489, 1162261467,
            3486784401, 10460353203, 31381059609, 94143178827, 282429536481,
            847288609443, 2541865828329, 7625597484987, 22876792454961, 68630377364883,
            205891132094649, 617673396283947, 1853020188851841};

        // ~ ceil( x * (2. / 3.)^y ) + 1
        inline size_t x_mul_pow2on3y(size_t x, size_t y) const
        {
            return (((x * 2) << y) / pow3[y] + 1) / 2;
        }

        size_t level_cap(size_t lv) const
        {
            size_t cap = k, d = compact.size() - lv - 1;
            while (d > 32)
            {
                cap = x_mul_pow2on3y(cap, 32);
                d -= 32;
            }
            return std::max(x_mul_pow2on3y(cap, d), static_cast<size_t>(8));
        }

        void update_size()
        {
            size = 0;
            for (const tvect &cpt : compact)
            {
                size += cpt.size();
            }
        }

        void update_max_size()
        {
            max_size = 0;
            for (size_t lv = 0; lv < compact.size(); ++lv)
            {
                max_size += level_cap(lv);
            }
        }

        void grow()
        {
            compact.push_back({});
            update_max_size();
        }

        void compress()
        {
            size_t maxlv = compact.size(), nextlv;
            for (size_t lv = 0; lv < maxlv; ++lv)
            {
                if (compact[lv].size() >= level_cap(lv))
                {
                    nextlv = lv + 1;
                    if (nextlv == maxlv)
                    {
                        grow();
                    }
                    // only level 0 is not sorted, O(k log k)
                    if (lv == 0)
                    {
                        std::sort(compact[0].begin(), compact[0].end());
                    }
                    work1.clear();
                    work2.clear();
                    // work1 is sorted
                    size_t w = static_cast<size_t>(1) << nextlv;
                    flip = !flip;
                    for (size_t i = flip ^ rbool(); i < compact[lv].size(); i += 2)
                    {
                        work1.push_back({compact[lv][i].v, w});
                    }
                    // merge O(k)
                    if (compact[nextlv].size())
                    {
                        std::merge(work1.begin(), work1.end(), compact[nextlv].begin(), compact[nextlv].end(), std::back_inserter(work2));
                        work2.swap(compact[nextlv]);
                    }
                    else
                    {
                        compact[nextlv].reserve(work1.size());
                        std::copy(work1.begin(), work1.end(), std::back_inserter(compact[nextlv]));
                    }
                    compact[lv].clear();
                    update_size();
                    // always perform lazy compaction
                    break;
                }
            }
        }

    public:
        /**
         * @brief Construct a new kll sketch object
         *
         * @param k max compactor size
         */
        kll_sketch(size_t k = 200)
            : k(k),
              work1(k),
              work2(k),
              work3(k * 32),
              work4(k * 32) // ~300-500kB work3/work4, avoid reallocation
        {
            grow();
        }

        /**
         * @brief Insert new value to sketch
         *
         * @param x
         */
        void insert(T x)
        {
            if (x > maxval)
            {
                maxval = x;
            }
            compact[0].push_back({x, 1ul});
            if (++size > max_size)
            {
                compress();
            }
        }

        /**
         * @brief Calculate density from sketch
         *
         * @tparam ValsIt value output iterator type
         * @tparam DensIt density output iterator type
         * @param[out] vals value output iterator
         * @param[out] dens density output iterator
         */
        template <class ValsIt, class DensIt>
        void density(ValsIt vals, DensIt dens, bool cumulative = false) const
        {
            work3.clear();
            work4.clear();
            work3.reserve(size);
            work4.reserve(size);

            size_t total_w = compact[0].size();

            // sort and merge level 0
            std::sort(compact[0].begin(), compact[0].end());
            std::copy(compact[0].begin(), compact[0].end(), std::back_inserter(work3));

            // merge other levels
            for (size_t lv = 1; lv < compact.size(); ++lv)
            {
                std::merge(work3.begin(), work3.end(), compact[lv].begin(), compact[lv].end(), std::back_inserter(work4));
                work3.swap(work4);
                work4.clear();
                // update weights
                total_w += (static_cast<size_t>(1) << lv) * compact[lv].size();
            }

            // unzip
            auto back_v = work3.front().v;
            auto back_w = work3.front().w;
            if (cumulative)
            {
                for (size_t i = 1; i < work3.size(); ++i)
                {
                    if (work3[i].v != back_v)
                    {
                        // insert new
                        *vals++ = back_v;
                        *dens++ = static_cast<P>(back_w) / total_w;
                        back_v = work3[i].v;
                    }
                    // accumulate weight
                    back_w += work3[i].w;
                }
            }
            else
            {
                for (size_t i = 1; i < work3.size(); ++i)
                {
                    if (work3[i].v != back_v)
                    {
                        // insert new
                        *vals++ = back_v;
                        *dens++ = static_cast<P>(back_w) / total_w;
                        back_v = work3[i].v;
                        back_w = work3[i].w;
                    }
                    else
                    {
                        // accumulate weight
                        back_w += work3[i].w;
                    }
                }
            }
            // insert last
            *vals++ = back_v;
            *dens++ = static_cast<P>(back_w) / total_w;
        }

        /**
         * @brief Calculate quantiles
         *
         * @tparam ProbsIt probs iterator type
         * @tparam ValsIt vals iterator type
         * @param[in] probs_first begin of probs input iterator
         * @param[in] probs_last end of probs input iterator
         * @param[out] vals value output iterator
         */
        template <class ProbsIt, class ValsIt>
        void quantile(ProbsIt probs_first, ProbsIt probs_last, ValsIt vals) const
        {
            // calculate cdf
            workv.clear();
            workp.clear();
            density(std::back_inserter(workv), std::back_inserter(workp), true);
            // find quantiles
            auto it = workp.begin();
            auto it_begin = workp.begin();
            auto it_end = workp.end();
            while (probs_first != probs_last)
            {
                it = std::lower_bound(it, it_end, *probs_first++);
                if (it == it_end)
                {
                    *vals++ = maxval;
                }
                else
                {
                    *vals++ = workv[it - it_begin];
                }
            }
        }

        void quantile(P prob, T &val) const
        {
            // calculate cdf
            workv.clear();
            workp.clear();
            density(std::back_inserter(workv), std::back_inserter(workp), true);
            auto it = std::lower_bound(workp.begin(), workp.end(), prob);
            if (it == workp.end())
            {
                val = maxval;
            }
            else
            {
                val = workv[it - workp.begin()];
            }
        }

        void merge(const kll_sketch &rhs)
        {
            while (compact.size() < rhs.compact.size())
            {
                grow();
            }
            // insert to level 0, not sorted
            compact[0].insert(compact[0].end(), rhs.compact[0].begin(), rhs.compact[0].end());
            // merge higher levels, sorted
            for (size_t lv = 1; lv < rhs.compact.size(); ++lv)
            {
                work1.clear();
                std::merge(compact[lv].begin(), compact[lv].end(), rhs.compact[lv].begin(), rhs.compact[lv].end(), std::back_inserter(work1));
                work1.swap(compact[lv]);
            }
            // compress if needed
            update_size();
            if (size > max_size)
            {
                compress();
            }
        }

        void reset()
        {
            compact.clear();
            grow();
        }
    };
}
